# -*- coding: utf-8 -*-

# @Time    : 2019/1/31 22:56
# @Author  : Chen
# @File    : training.py
# @Software: PyCharm

import os
import numpy as np
import tensorflow as tf
import input_data
import model
from PIL import Image
import matplotlib.pyplot as plt

def training():
    # 变量声明
    N_CLASSES = 2  # 猫和狗
    IMG_W = 208  # resize图像，太大的话训练时间久
    IMG_H = 208
    BATCH_SIZE = 16
    CAPACITY = 2000
    MAX_STEP = 10000  # 一般大于10K
    learning_rate = 0.0001  # 一般小于0.0001

    # 获取批次batch
    train_dir = '..\\data\\dogs-vs-cats\\train\\'
    logs_train_dir = '..\\data\\dogs-vs-cats\\logs\\train\\'
    train, train_label = input_data.get_files(train_dir)
    train_batch, train_label_batch = input_data.get_batch(train,
                                                          train_label,
                                                          IMG_W,
                                                          IMG_H,
                                                          BATCH_SIZE,
                                                          CAPACITY)

    # 操作定义
    train_logits = model.inference(train_batch, BATCH_SIZE, N_CLASSES)
    train_loss = model.losses(train_logits, train_label_batch)
    train_op = model.trainning(train_loss, learning_rate)
    train_acc = model.evaluation(train_logits, train_label_batch)
    summary_op = tf.summary.merge_all() # 这个是log汇总记录
    # 产生一个会话
    sess = tf.Session()
    # 产生一个writer来写log文件
    train_writer = tf.summary.FileWriter(logs_train_dir, sess.graph)
    # 产生一个saver来存储训练好的模型
    saver = tf.train.Saver()
    # 所有节点初始化
    sess.run(tf.global_variables_initializer())
    # 队列监控
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess = sess, coord = coord)

    # 进行batch的训练
    try:
        # 执行MAX_STEP步的训练，一步一个batch
        for step in np.arange(MAX_STEP):
            if coord.should_stop():
                break
            # 启动以下操作节点，有个疑问，为什么train_logits在这里没有开启？
            _, tra_loss, tra_acc = sess.run([train_op, train_loss, train_acc])
            # 每隔50步打印一次当前的loss以及acc，同时记录log，写入writer
            if step % 50 == 0:
                print('Step %d, train loss = %.2f, train accuracy = %.2f%%' %(step, tra_loss, tra_acc*100.0))
                summary_str = sess.run(summary_op)
                train_writer.add_summary(summary_str, step)
            # 每隔2000步，保存一次训练好的模型
            if step % 2000 == 0 or (step + 1) == MAX_STEP:
                checkpoint_path = os.path.join(logs_train_dir, 'model.ckpt')
                saver.save(sess, checkpoint_path, global_step = step)
    except tf.errors.OutOfRangeError:
        print('Done training -- epoch limit reached')
    finally:
        coord.request_stop()
    coord.join(threads = threads)
    sess.close()

# 测试一张图片
def get_one_image(train): # 获取一张图片
    # 输入参数：train,训练图片的路径
    # 返回参数：image，从训练图片中随机抽取一张图片

    n = len(train)
    ind = np.random.randint(0, n)
    img_dir = train[ind]

    image = Image.open(img_dir)
    plt.imshow(image)
    image = image.resize([208, 208])
    image = np.array(image)
    return  image
def evaluate_one_image(image_array): # 测试图片
    with tf.Graph().as_default():
        BATCH_SIZE = 1
        N_CLASSES = 2

        image = tf.cast(image_array, tf.float32)
        image = tf.image.per_image_standardization(image)
        image = tf.reshape(image, [1, 208, 208, 3])

        logit = model.inference(image, BATCH_SIZE, N_CLASSES)

        logit = tf.nn.softmax(logit)

        x = tf.placeholder(tf.float32, shape = [208, 208, 3])

        logs_train_dir = '..\\data\\dogs-vs-cats\\logs\\train\\'

        saver = tf.train.Saver()

        with tf.Session() as sess:
            print('Reading checkpoints...')
            ckpt = tf.train.get_checkpoint_state(logs_train_dir)
            if ckpt and ckpt.model_checkpoint_path:
                global_step = ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1]
                saver.restore(sess, ckpt.model_checkpoint_path)
                print('Loading success, global_step is %s' % global_step)
            else:
                print('No checkpoint file found')
            prediction = sess.run(logit, feed_dict={x: image_array})
            max_index = np.argmax(prediction)
            if max_index == 0:
                print('This is a cat with possibility %.6f' % prediction[:, 0])
            else:
                print('This is a dog with possibility %.6f' % prediction[:, 1])